package com.xp.ai.util;

import com.xp.ai.conf.Constants;
import dev.langchain4j.community.model.zhipu.ZhipuAiChatModel;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.chat.request.ResponseFormat;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.moderation.ModerationModel;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.model.ollama.OllamaStreamingChatModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.openai.OpenAiLanguageModel;
import dev.langchain4j.model.openai.OpenAiModerationModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;

import java.time.Duration;

public class ModelUtils {


    /**
     * 构建最简单，普通的聊天大模型
     * @return ChatLanguageModel
     */
    public static ChatLanguageModel getOllamaNormalModel(){
        return OllamaChatModel.builder()
                .baseUrl(Constants.OLLAMA_BASE_URL)
                .modelName(Constants.OLLAMA_MODEL_NAME)
                .build();
    }



    public static ChatLanguageModel getOllamaJsonModel(){
        return OllamaChatModel.builder()
                .baseUrl(Constants.OLLAMA_BASE_URL)
                .modelName(Constants.OLLAMA_MODEL_NAME)
                .responseFormat(ResponseFormat.JSON)
                .build();
    }


    /**
     * 获取流式聊天大模型
     * @return ChatLanguageModel
     */
    public static StreamingChatLanguageModel getOllamaStreamModel(){
        return OllamaStreamingChatModel.builder()
                .baseUrl(Constants.OLLAMA_BASE_URL)
                .modelName(Constants.OLLAMA_MODEL_NAME)
                .temperature(0.0)
                .build();
    }


    /***
     * 用OpenAi 兼容模式获取到火山的DeepSeekR1Net大模型
     * @return
     */
    public static ChatLanguageModel  getHuoshanR1Model(){
        return OpenAiChatModel.builder()
                .baseUrl(ApiKey.HUOSHAN_BASE_URL)
                .modelName(ApiKey.HUOSHAN_R1_ID)
                .apiKey(ApiKey.HUOSHAN_R1_KEY)
                .build();
    }


    /***
     * 获取火山大模型 deepseek v3 大模型
     * @return
     */
    public static ChatLanguageModel   getHuoshanv3Model(){
        return OpenAiChatModel.builder()
                .baseUrl(ApiKey.HUOSHAN_BASE_URL)
                .modelName(ApiKey.HUOSHAN_V3_ID)
                .apiKey(ApiKey.HUOSHAN_R1_KEY)
                .build();
    }


    /***
     * 获取火山大模型 deepseek v3 大模型
     * @return
     */
    public static ChatLanguageModel   getDeepSeeV3Model(){
        return OpenAiChatModel.builder()
                .baseUrl(ApiKey.DEEP_SEEK_BASE_URL)
                .modelName(ApiKey.DEEP_SEEK_V3_ID)
                .apiKey(ApiKey.DEEP_SEEK_API_KEY)
                .build();
    }


    public static ChatLanguageModel   getAliZpModel(){
        return ZhipuAiChatModel.builder()
                .baseUrl(ApiKey.ALI_ZP_BASE_URL)
                .model(ApiKey.ALI_ZP_MODEL_ID)
                .apiKey(ApiKey.ALI_ZP_API_KEY)
                .callTimeout(Duration.ofSeconds(10L))
                .connectTimeout(Duration.ofSeconds(10L))
                .writeTimeout(Duration.ofSeconds(10L))
                .readTimeout(Duration.ofSeconds(10L))
                .maxToken(8196)
                .build();
    }

    public static ChatLanguageModel   getAliQwModel(){
        return OpenAiChatModel.builder()
                .baseUrl(ApiKey.ALI_ZP_BASE_URL)
                .modelName(ApiKey.ALI_QW32_MODEL_ID)
                .apiKey(ApiKey.ALI_ZP_API_KEY)
                .build();
    }

    public static EmbeddingModel getGjEmbeddingModel(){
        return OpenAiEmbeddingModel.builder()
                .baseUrl(ApiKey.GJ_BASE_URL)
                .apiKey(ApiKey.GJ_API_KEY)
                .modelName(ApiKey.GJ_EMBEDDING_MODEL)
                .build();
    }


    public static EmbeddingStore<TextSegment>  getPgEmbeddingStore(){
        return PgVectorEmbeddingStore.builder()
                .host("localhost")
                .port(5433)
                .user("root")
                .password("123456")
                .database("postgres")
                .table("question_embedding")
                .dimension(1024)
                .useIndex(true)
                .indexListSize(100)
                .createTable(true)
                .dropTableFirst(false)
                .build();
    }
}
